
using System.Collections.Concurrent;
using System.Linq;
using System.Net;
using System.Net.Sockets;

namespace Dong
{
    public class RemoteClient
    {
        public IPEndPoint PublicEndPoint { get; }
        public IPEndPoint PrivateEndPoint { get; }
        public RemoteClient(IPEndPoint publicEndPoint, IPEndPoint privateEndPoint)
        {
            PublicEndPoint = publicEndPoint;
            PrivateEndPoint = privateEndPoint;
        }
    }

    public class Server
    {
        private readonly Socket _socket;
        public Server(int port)
        {
            _socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
            _socket.Bind(new IPEndPoint(IPAddress.Any, port));
        }

        public void Start()
        {
            _socket.Listen(5);
            while (true)
            {
                var socket = _socket.Accept();
                DoWork(socket);
            }
        }

        private ConcurrentDictionary<RemoteClient, Socket> _clients = new();

        private async void DoWork(Socket socket)
        {
            await Task.Run(() =>
            {
                // Verify socket
                const int checkBufferLength = 3;
                byte[] checkValues = new byte[checkBufferLength];
                var random = new Random(DateTime.Now.Millisecond);
                random.NextBytes(checkValues);
                socket.Send(checkValues);
                byte[] checkResult = new byte[checkBufferLength];
                ReceiveToFullBytes(socket, checkResult, 0, checkResult.Length);
                for (int i = 0; i < checkBufferLength; ++i)
                {
                    if (checkValues[i] + checkResult[i] != byte.MaxValue)
                    {
                        socket.Disconnect(false);
                        socket.Close();
                        socket.Dispose();
                    }
                }
                // Record the remote IP
                byte[] buffer = new byte[12];
                ReceiveToFullBytes(socket, buffer, 0, buffer.Length);
                int port = BitConverter.ToInt32(buffer, 8);
                var client = new RemoteClient((IPEndPoint)socket.RemoteEndPoint!, new IPEndPoint(IPAddress.Parse($"{buffer[0]}.{buffer[1]}.{buffer[2]}.{buffer[3]}"), port));
                _clients.TryAdd(client, socket);
                while (true)
                {
                    ReceiveToFullBytes(socket, buffer, 0, 4);
                    int commandId = BitConverter.ToInt32(buffer, 0);
                    switch (commandId)
                    {
                        case 0:
                            int clientCount = _clients.Count - 1;
                            socket.Send(BitConverter.GetBytes(clientCount));
                            foreach (var c in _clients.Keys)
                            {
                                if (c.PublicEndPoint.Equals(socket.RemoteEndPoint))
                                    continue;
                                socket.Send(BitConverter.GetBytes(c.PublicEndPoint.Address.Address));
                                socket.Send(BitConverter.GetBytes(c.PublicEndPoint.Port));
                                socket.Send(BitConverter.GetBytes(c.PrivateEndPoint.Address.Address));
                                socket.Send(BitConverter.GetBytes(c.PrivateEndPoint.Port));
                            }
                            break;
                        case 1:
                            ReceiveToFullBytes(socket, buffer, 0, buffer.Length);
                            long address = BitConverter.ToInt64(buffer, 0);
                            port = BitConverter.ToInt32(buffer, 8);
                            var targetClient = _clients.Keys.FirstOrDefault(c => c.PublicEndPoint.Address.Address == address && c.PublicEndPoint.Port == port);
                            if (targetClient != null)
                            {
                                var s = _clients[targetClient];
                                s.Send(BitConverter.GetBytes(client.PublicEndPoint.Address.Address));
                                s.Send(BitConverter.GetBytes(client.PublicEndPoint.Port));
                                s.Send(BitConverter.GetBytes(client.PrivateEndPoint.Address.Address));
                                s.Send(BitConverter.GetBytes(client.PrivateEndPoint.Port));
                                socket.Send(BitConverter.GetBytes(targetClient.PublicEndPoint.Address.Address));
                                socket.Send(BitConverter.GetBytes(targetClient.PublicEndPoint.Port));
                                socket.Send(BitConverter.GetBytes(targetClient.PrivateEndPoint.Address.Address));
                                socket.Send(BitConverter.GetBytes(targetClient.PrivateEndPoint.Port));
                            }
                            break;
                    }
                }
            });
        }

        private static void ReceiveToFullBytes(Socket socket, byte[] buffer, int offset, int length)
        {
            int receivedCount = 0;
            while (receivedCount < length)
            {
                receivedCount += socket.Receive(buffer, offset + receivedCount, length - receivedCount, SocketFlags.None);
            }
        }
    }
}